from greedy_policy import TC_epsilon_greedy_policy,NC_epsilon_greedy_policy
from visit_maintain import visit_main
# from numba import jit

def choose_action(S,IU,NIU,MTN,flight_state,aircraft_maintain,base_cap,base, aircraft_action,aircraft,aircraft_time,aircraft_flight_innumber,states,final_arr_time,all_cap,reward_sample,epsilon,
                  qfun,action_result):
    IU = IU
    NIU = NIU
    MTN = MTN
    S = S
    flight_state = flight_state
    aircraft_maintain = aircraft_maintain
    base_cap = base_cap
    base = base
    aircraft_action =  aircraft_action
    aircraft = aircraft
    aircraft_time = aircraft_time
    aircraft_flight_innumber = aircraft_flight_innumber
    states = states
    final_arr_time = final_arr_time
    all_cap = all_cap
    reward_sample = reward_sample
    epsilon = epsilon
    qfun = qfun
    action_result = action_result
    NC = []
    TC = []
    # CMAX = 15  # 两次维修间最大允许起降架次
    # TMAX = 2400  # 两次维修间最大允许累计飞行时间
    #STEP A
    # 此时没有在用飞机时，从未使用的飞机中进行选择，并且要考虑维修完成的飞机序列状况
    if len(IU)==0:
        action = min(aircraft_action.index)  # 此时筛选得出可用飞机的索引号，即课代表选择的飞机号
        aircraft[action].append(S) #将飞机飞行状态计入集合列表
        aircraft_time[action] += flight_state.at[S, 'flight_time']  #记录累计飞行时间
        aircraft_flight_innumber[action] += 1 #记录累计起降架次
        IU.append(action)
        NIU.remove(action)
        states.remove(S)#删除所进行分配过动作的状态
        action_result[S] = action #将每个状态的动作汇聚到一个集合中进行直观展示
        return states

    #STEP B  IU列表不为空的时候可以判断是否存在连接情况
    if not len(IU) == 0:
        for i in IU:
          if (aircraft_flight_innumber[i] == 14):
             time = []
             for j in aircraft[i]:  # 飞机i执行的航班集合信息
                 time.append(flight_state.at[j, 'arr_time'])
             final_time = max(time)
             final_arr_time[i] = final_time
             final_time_index = time.index(final_time)  # 找到飞机状态最晚到达时间与最晚时间对应索引位置
             final_air = flight_state.at[aircraft[i][final_time_index], 'arr_air']  # 筛选最后一趟航班对应的到达机场
             maintain_num, punish = visit_main(aircraft_time[i], final_time, final_air,base_cap, base, all_cap, aircraft[i][final_time_index], states)
             aircraft_maintain[i] = maintain_num  # 将此时的维修航段存入飞机维修集合信息中
             action_result[maintain_num] = i  # 将飞机号分配给维修航段，也是满足状态取飞机号的形式
             MTN.append(i)  # 将其纳入维修列表进行维修约束,确定最终维修弧
             IU.remove(i)
             aircraft_time[i] = 0  # 清零累计飞行时间
             aircraft_flight_innumber[i] = 0  # 清零起降架次
             maintain_air = flight_state.at[maintain_num, 'arr_air']
             if maintain_num in states:
                 states.remove(maintain_num)
             if not maintain_num == aircraft[i][final_time_index]:  # 不是就地维修，那一定是不惩罚的情况
                 aircraft[i].append(maintain_num)
                 base_cap[base.index(maintain_air)] += 1  # 维修机场的在用维修容量+1
                 if (flight_state.at[maintain_num,'dep_time']-flight_state.at[aircraft[i][final_time_index],'arr_time'])<=90 and (flight_state.at[maintain_num, 'dep_time'] - flight_state.at[aircraft[i][final_time_index], 'arr_time']) >= 45:
                     reward_sample[maintain_num] = 10
             if maintain_num == aircraft[i][final_time_index]:
                if punish == True:
                     reward_sample[maintain_num] = -10
        S = states[0]
        #得出状态s的出发时间，以判断是否可以进行连接
        dep_air = flight_state.at[S, 'dep_air']
        dep_time = flight_state.at[S, 'dep_time']
        for i in IU: #判断在用的飞机是否满足NC连接约束，如果满足则将其挪至nc连接列表
            time = []
            for j in aircraft[i]: #飞机i执行的航班集合信息
                time.append(flight_state.at[j, 'arr_time'])
            final_time = max(time)
            final_arr_time[i] = final_time
            final_time_index = time.index(final_time) #找到飞机状态最晚到达时间与最晚时间对应索引位置
            final_air = flight_state.at[aircraft[i][final_time_index], 'arr_air'] #筛选最后一趟航班对应的到达机场
            #找到所有在用飞机中可以与此状态形成NC连接的所有NC情况
            if ((aircraft_flight_innumber[i])<14 and (final_air == dep_air) and (dep_time-final_time)>=45 and (flight_state.at[S, 'flight_time']+aircraft_time[i])<=2400 ): #满足时间地点机场衔接约束
                NC.append(i) #查找满足衔接条件的飞机加入到NC列表中

            #对于飞机进行维修检查约束限制的考虑，即飞机i不可以执行此时S的状态，为其确定维修航段，并清空飞机状态、可用状态以及增加机场容量的目前所用信息
            if ( (aircraft_flight_innumber[i] == 14) or (flight_state.at[S, 'flight_time']+aircraft_time[i]>2400)  ):
               maintain_num , punish = visit_main(aircraft_time[i],final_time,final_air,base_cap,base,all_cap,aircraft[i][final_time_index],states)  #-----------晚点检查----------
               aircraft_maintain[i] = maintain_num  # 将此时的维修航段存入飞机维修集合信息中
               action_result[maintain_num] = i  # 将飞机号分配给维修航段，也是满足状态取飞机号的形式
               MTN.append(i)  # 将其纳入维修列表进行维修约束,确定最终维修弧
               IU.remove(i)
               aircraft_time[i] = 0  # 清零累计飞行时间
               aircraft_flight_innumber[i] = 0  # 清零起降架次
               maintain_air = flight_state.at[maintain_num, 'arr_air']
               if maintain_num in states:
                   states.remove(maintain_num)
               if not maintain_num==aircraft[i][final_time_index]: #不是就地维修，那一定是不惩罚的情况
                   aircraft[i].append(maintain_num)
                   base_cap[base.index(maintain_air)] += 1  # 维修机场的在用维修容量+1
                   if (flight_state.at[maintain_num,'dep_time']-flight_state.at[aircraft[i][final_time_index],'arr_time'])<=90 and (flight_state.at[maintain_num,'dep_time']-flight_state.at[aircraft[i][final_time_index],'arr_time'])>=45:
                       reward_sample[maintain_num] = 10
               else:
                   if punish == True:
                       reward_sample[maintain_num] = -10

        if len(NC)==0:
            check_time = flight_state.at[S, 'dep_time']
            if not len(MTN) == 0:  # 当MTN不为空时
                for i in MTN:
                    # 查找i对应的维修航段
                    M_S = aircraft_maintain[i]  # 找到此时飞机i所在的维修航段及其状况
                    M_start = flight_state.at[M_S, 'arr_time']
                    M_air = flight_state.at[M_S, 'arr_air']
                    M_end = M_start + 480  # 该飞机结束维修的时间
                    if ((M_end < check_time) and (M_air==flight_state.at[S, 'dep_air'])):
                        MTN.remove(i)
                        NIU.append(i)
            s_airport = flight_state.at[S, 'dep_air']  # 判断此航班的出发机场是哪里，依次选取未用飞机基地机场是此出发机场的集合
            NIU1=NIU
            for i in NIU:
                if not len(aircraft[i])==0:
                    if not (flight_state.at[aircraft[i][-1], 'arr_air']==s_airport):
                        NIU1.remove(i)
            # aircraft_available = aircraft_action[aircraft_action.index.isin(NIU1)]  # 将NIU中飞机的信息筛选出来
            aircraft_available = aircraft_action[[x in NIU1 for x in aircraft_action.index]]
            # s_action = aircraft_available
            action = min(aircraft_available.index)  # 此时筛选得出可用飞机的索引号，即课代表选择的飞机号
            aircraft[action].append(S)  # 将飞机飞行状态计入集合列表
            aircraft_time[action] += flight_state.at[S, 'flight_time']  # 记录累计飞行时间
            aircraft_flight_innumber[action] += 1  # 记录累计起降架次
            IU.append(action)
            NIU.remove(action)
            states.remove(S)  # 删除所进行分配过动作的状态
            action_result[S] = action  # 将每个状态的动作汇聚到一个集合中进行直观展示
            return states
        #STEP C
        if not len(NC)==0:
            #检查NC连接中是否存在TC连接，即衔接时间在45min到1.5个小时之间
            for i in NC:
                if (dep_time -final_arr_time[i])<=90 and (dep_time -final_arr_time[i])>=45:
                    TC.append(i)
                    NC.remove(i)#完成飞机号从NC往TC的移动
            #STEP D
            if not len(TC)==0:#TC不为空的条件下使用贪婪策略进行路径的选择
                final_action = TC_epsilon_greedy_policy(S, TC, epsilon, qfun) #EPILON与qfun 还没定义，后面补充
                aircraft_time[final_action] += flight_state.loc[S, 'flight_time']  # 记录累计飞行时间
                aircraft_flight_innumber[final_action] += 1  # 记录累计起降架次
                reward_sample[S] = 10
                if final_action in NIU:
                   NIU.remove(final_action)
                if not final_action in IU:
                   IU.append(final_action)
                states.remove(S)  # 删除所进行分配过动作的状态
                action_result[S] = final_action  # 将每个状态的动作汇聚到一个集合中进行直观展示
                aircraft[final_action].append(S)  # 将飞机飞行状态计入集合列表
                return states
            #STEP E
            else:
                #此函数对于维修状态、机场分布、维修集合等信息进行一个简单的更新，可以不输出，后续操作可以直接使用
                final_action = NC_epsilon_greedy_policy(S, NC, NIU,MTN,aircraft_maintain,epsilon, qfun,aircraft) #EPILON与qfun 还没定义，后面补充
                aircraft_time[final_action] += flight_state.at[S, 'flight_time']  # 记录累计飞行时间
                aircraft_flight_innumber[final_action] += 1  # 记录累计起降架次
                if final_action in NIU:
                    NIU.remove(final_action)
                if not final_action in IU:
                    IU.append(final_action)
                states.remove(S)  # 删除所进行分配过动作的状态
                action_result[S] = final_action  # 将每个状态的动作汇聚到一个集合中进行直观展示
                aircraft[final_action].append(S)  # 将飞机飞行状态计入集合列表
                return states














